# 实现指定频率的滤波
import numpy.fft as nf
import numpy as np
import matplotlib.pyplot as plt
import os
import linecache
import re

# 设置滤波的频率和带宽
freqs = [0, 50, 150]   # 需要过滤的频率值
bandwidth = 4   # 频率带宽

# 读取文件
folder = 'DataToBeFiltered'  # 数据存储的文件夹名
files = os.listdir(folder)
for file in files:
    if file.startswith(".") or file.endswith("1") or file.endswith(".png") or file.endswith(".pdf"):  #
        # 过滤掉以“.”开头的文件和以“1”、“.png”、".pdf"结尾的文件
        continue

    scale = float(re.split(r'[\t\n]', linecache.getline(folder + r'/' + file, 2))[1])  # 读取电压信号的比例值
    data = np.loadtxt(file, skiprows=9)    # 跳过前9行数据
    t = data[:, 0]      # 第一列数据为时间序列
    f_t = data[:, 1]    # 第二列数据为函数序列

    fig, axes = plt.subplots(2, 2, sharex='col', sharey='col')  # 画两个子图，分别为时域图和频域图
    fig.dpi = 300   # 设置画出的图像的dpi，dpi越大，图像越清晰
    # 打印波形图
    axes[0, 0].plot(t, f_t)
    axes[0, 0].set_ylabel("signal")

    # 傅立叶变换
    F = nf.fft(f_t)
    N = f_t.size    # 采样点数
    delta_t = t[1] - t[0]     # 采样间隔
    f = nf.fftfreq(f_t.size, delta_t)   # 频域的横坐标序列
    mask = np.where(f > 0)  # 只分析频率大于零的部分（小于零的部分与之完全对称）
    A = 2*abs(F[mask]) / N  # 频谱图的纵坐标变换为实际振幅
    # 打印频谱图
    axes[0, 1].plot(f[mask], A)
    axes[0, 1].set_ylabel("$A$")

    for freq in freqs:
        lower_freq = freq - bandwidth / 2
        upper_freq = freq + bandwidth / 2
        F_f = F * ((abs(f) < lower_freq) + (abs(f) > upper_freq))    # 滤波之后的傅立叶变换值
        f_t_f = nf.ifft(F_f)  # 傅立叶逆变换
        f_t = f_t_f.real
        F = F_f
    axes[1, 0].plot(t, f_t_f.real, color='red')  # 画出滤波之后的时域波形
    axes[1, 0].set_xlabel("time (s)")
    axes[1, 0].set_ylabel("signal")

    # 打印滤波之后的频谱图
    A_f = 2*abs(F_f[mask]) / N  # 频谱图的纵坐标变换为实际振幅
    axes[1, 1].plot(f[mask], A_f, color='red')
    axes[1, 1].set_xlabel("frequency (Hz)")
    axes[1, 1].set_ylabel("$A$")

    # 分析并输出实验结果
    meas_freq = float(file.split('Hz')[0])     # 从文件名中读取激振器的频率值
    sign_freq = (f[f > 0])[np.argmax(abs(F[mask]))]       # 找到频谱图山振幅最大时对应的频率值
    amplitude = np.max(A_f)   # 最大振幅值
    print(f"分量最大的频率分量为：{sign_freq:.2f} Hz, 对应的幅值为：{amplitude:.2f} V")
    devi_freq = sign_freq - meas_freq   # 绝对误差
    real_devi = devi_freq / meas_freq   # 相对误差
    print(f"偏差为：{devi_freq:.2f} Hz（{real_devi:.2%}）".format(devi_freq, real_devi))
    text = f"f = {sign_freq:.2f} Hz, A = {amplitude:.2f} V, error = {devi_freq:.2f} Hz, {real_devi:.2%}"
    plt.suptitle(text)  # 将误差作为图的标题，显示在图的上方

    plt.tight_layout(rect=[0, 0, 1, 0.92])  # 设置子图布局，避免重叠
    plt.savefig(folder + r'/' + file + '.png')    # 将分析结果保存在以频率信息命名的png图片中
plt.show()
